import os
from os.path import join, isdir, isfile
import numpy as np


def merge_navigate_touchpoint(touchpoint_root, navigate_root, save_dir, seq_name, sampling_rate_touchpoint=1, sampling_rate_navigate=1):
    os.makedirs(join(save_dir, seq_name, "h1_kinematic_motions"), exist_ok=True)
    os.system("cp -rp {} {}".format(join(touchpoint_root, seq_name, "scene_mesh.obj"), join(save_dir, seq_name, "scene_mesh.obj")))

    for fn in os.listdir(join(touchpoint_root, seq_name, "h1_kinematic_motions")):
        touchpoint_fp = join(touchpoint_root, seq_name, "h1_kinematic_motions", fn)
        navigate_fp = join(navigate_root, seq_name, "h1_kinematic_motions", fn.replace("_data.npz", "_demo_motion_0.npz"))
        assert isfile(navigate_fp), navigate_fp
        touchpoint_motion = np.load(touchpoint_fp, allow_pickle=True)["arr_0"].item()
        navigate_motion = np.load(navigate_fp, allow_pickle=True)["arr_0"].item()
        merged_motion = {}
        for key in touchpoint_motion:
            if isinstance(touchpoint_motion[key], dict):
                merged_motion[key] = {}
                for link_name in touchpoint_motion[key].keys():
                    # TODO: add interpolation
                    merged_motion[key][link_name] = np.concatenate((navigate_motion[key][link_name][::sampling_rate_navigate], touchpoint_motion[key][link_name][::sampling_rate_touchpoint]), axis=0)
            else:
                merged_motion[key] = np.concatenate((navigate_motion[key][::sampling_rate_navigate], touchpoint_motion[key][::sampling_rate_touchpoint]), axis=0)
        np.savez(join(save_dir, seq_name, "h1_kinematic_motions", fn), merged_motion)


if __name__ == "__main__":
    touchpoint_root = "./CORE4D_retargeted_data_touchpoint"
    navigate_root = "./CORE4D_walks_UniHSI_copyrot_0920"
    save_dir = "./CORE4D_retargeted_data_walktouchpoint"

    os.makedirs(save_dir, exist_ok=True)

    for seq_name in os.listdir(touchpoint_root):
        merge_navigate_touchpoint(touchpoint_root, navigate_root, save_dir, seq_name, sampling_rate_touchpoint=1, sampling_rate_navigate=2)
